﻿/*
* Copyright (c) 2012 Kyle Hayward
* 
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
* 
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
* 
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
* THE SOFTWARE.
*/

using System;
using System.Runtime.InteropServices;
using System.Collections.Generic;
using System.Linq;
using System.IO;
using System.Windows.Forms;

using SlimDX;
using SlimDX.Windows;
using SlimDX.RawInput;

using OptixDotNet;
using OptixDotNetUtils;
using OptixMath = OptixDotNet.Math.MathUtils;

using Buffer = OptixDotNet.Buffer;
using BoundingBox = OptixDotNet.Math.BoundingBox;
using OptixV3 = OptixDotNet.Math.Vector3;

namespace InstantRadiosity
{
    static class App
    {
        /// <summary>
        /// The main entry point for the application.
        /// </summary>
        [STAThread]
        static void Main( string[] args )
        {
            try
            {
                InstantRadiosity sample = new InstantRadiosity();
                sample.Run();
            }
            catch ( Exception ex )
            {
                Console.WriteLine( ex.Message );
                throw;
            }
        }
    }

    public struct OptixLight
    {
        public Vector3 position;
        public Vector3 color;
        public Vector3 direction;
        public float intensity;
    }

    public class InstantRadiosity
    {
        public static readonly float MaxVoxelSize = 6.0f;
        public static readonly uint SqrtNumVPLs = 32;
        public static readonly uint NumBounces = 3;
        public static readonly uint TotalVPLs = SqrtNumVPLs * SqrtNumVPLs * NumBounces;
        public static readonly uint SqrtNumBounceVPLs = SqrtNumVPLs * NumBounces;

        RenderForm mForm;
        Renderer mRenderer;

        Context mContext;
        D3DBuffer mOutputBuffer;

        List<BoundingBox> mVoxels = new List<BoundingBox>();
        int mFirstLeafIndex = -1;

        D3DCamera mCamera;
        MouseInputEventArgs mPrevMouse;
        bool mLeftMouseDown = false;
        bool mRightMouseDown = false;
        bool mMiddleMouseDown = false;

        float mTimeDelta = 0.0f;

        public InstantRadiosity()
        {
            OptixDotNet.Utils.Time.InitTime();

            mForm = new RenderForm( "InstantRadiosity" );
            mForm.ClientSize = new System.Drawing.Size( 1280, 720 );

            mRenderer = new Renderer( mForm );
            SlimDX.RawInput.Device.RegisterDevice( SlimDX.Multimedia.UsagePage.Generic, SlimDX.Multimedia.UsageId.Mouse, DeviceFlags.None );
            SlimDX.RawInput.Device.RegisterDevice( SlimDX.Multimedia.UsagePage.Generic, SlimDX.Multimedia.UsageId.Keyboard, DeviceFlags.None );
            SlimDX.RawInput.Device.MouseInput += new EventHandler<MouseInputEventArgs>( MouseInput );
            SlimDX.RawInput.Device.KeyboardInput += new EventHandler<KeyboardInputEventArgs>( KeyboardInput );
            mPrevMouse = new MouseInputEventArgs( MouseMode.AbsoluteMovement, MouseButtonFlags.None, 0, 0, 0, 0, 0, IntPtr.Zero );

            AreaLight light = new AreaLight();
            light.Position = new Vector4( 0.0f, 17.0165f, 0.0f, 25.0f );
            light.Color = new Vector4( 1.0f, 1.0f, 0.633333f, 1.0f ) * 10.0f;
            light.Corner = new Vector4( 2.8096f, 17.1165f, 2.3659f, 0.0f );
            light.V1 = new Vector4( -5.0f, 0.0f, 0.0f, 0.0f );
            light.V2 = new Vector4( 0.0f, 0.0f, 5.0f, 0.0f );
            mRenderer.AddLight( ref light );

            InitializeOptix( light );
            LoadContent();
        }

        public void InitializeOptix( AreaLight light )
        {
            Console.WriteLine( "Initializing Optix scene" );

            string rayGenPath = Path.GetFullPath( "../../Assets/Cuda/InstantRadiosity.ptx" );

            /*-----------------------------------------------
             * Create the Optix context
             *-----------------------------------------------*/
            mContext = new Context();
            mContext.RayTypeCount = 1;
            mContext.EntryPointCount = 1;
            mContext.D3D11Device = mRenderer.Device.ComPointer;

            mContext[ "scene_epsilon" ].Set( 0.001f );
            mContext[ "MaxBounces" ].Set( NumBounces );
            mContext[ "radiance_ray_type" ].Set( 0u );
            mContext[ "IndirectIntensity" ].Set( 1.0f );

            /*-----------------------------------------------
             * Create the ray-generation and exception programs
             *-----------------------------------------------*/
            Program rayGen = new Program( mContext, rayGenPath, "instant_radiosity" );
            Program exception = new Program( mContext, rayGenPath, "exception" );
            Program miss = new Program( mContext, rayGenPath, "miss" );
            miss[ "bg_color" ].Set( 100 / 255.0f, 149 / 255.0f, 237 / 255.0f );

            mContext.SetRayGenerationProgram( 0, rayGen );
            mContext.SetExceptionProgram( 0, exception );
            mContext.SetRayMissProgram( 0, miss );

            /*-----------------------------------------------
             * Create the VPL output buffer:
             * 
             * the size of this buffer is the number of photons emited from the light * the number of bounces
             *-----------------------------------------------*/
            mOutputBuffer = new D3DBuffer( mContext, new D3DBufferDesc()
            {
                Api = D3DApi.Direct3D11,
                Width = TotalVPLs,
                ElemSize = (uint)PointLight.SizeInBytes,
                Format = Format.User,
                Type = BufferType.Output,
                Resource = mRenderer.VPLBuffer.ComPointer
            } );
            mContext[ "output_vpls" ].Set( mOutputBuffer );

            /*-----------------------------------------------
             * Create Optix light
             *-----------------------------------------------*/
            OptixLight optixLight = new OptixLight();
            optixLight.position = new Vector3( light.Position.X, light.Position.Y, light.Position.Z );
            optixLight.color = new Vector3( light.Color.X, light.Color.Y, light.Color.Z );
            optixLight.direction = new Vector3( 0.0f, -1.0f, 0.0f );
            optixLight.intensity = 1.0f;

            mContext[ "Light" ].SetUserData<OptixLight>( optixLight );

            /*-----------------------------------------------
             * Create seed buffer
             *-----------------------------------------------*/
            BufferDesc noiseDesc = new BufferDesc() { Width = SqrtNumVPLs, Height = SqrtNumVPLs, Format = Format.UInt2, Type = BufferType.Input };
            Buffer noiseBuffer = new Buffer( mContext, noiseDesc );

            Random rand = new Random();
            BufferStream stream = noiseBuffer.Map();
            for ( int i = 0; i < SqrtNumVPLs * SqrtNumVPLs; i++ )
            {
                OptixDotNet.Math.UInt2 seed = new OptixDotNet.Math.UInt2( (uint)rand.Next(), (uint)rand.Next() );
                stream.Write<OptixDotNet.Math.UInt2>( seed );
            }
            noiseBuffer.Unmap();

            mContext[ "seed_buffer" ].Set( noiseBuffer );
        }

        public void LoadContent()
        {
            string shaderPath = Path.GetFullPath( "../../Assets/Cuda/InstantRadiosity.ptx" );

            Shader shader = mRenderer.CreateShader( "../../Assets/Shaders/Phong.hlsl", "VSMain", "PSMain" );

            /*-----------------------------------------------
             * Create the material that will be executed when there is an intersection
             *-----------------------------------------------*/
            Material material = new Material( mContext );
            material.Programs[ 0 ] = new SurfaceProgram( mContext, RayHitType.Closest, shaderPath, "closest_hit_radiosity" );

            /*-----------------------------------------------
             * Load the geometry
             *-----------------------------------------------*/
            string intersectPath = Path.GetFullPath( "../../Assets/Cuda/triangle_mesh_interleaved.ptx" );

            D3DOBJLoader model = new D3DOBJLoader( "../../Assets/Models/cornell-dragon-exploded.obj", mRenderer );
            model.Shader = shader;

            model.Context = mContext;
            model.GeoGroup = new GeometryGroup( mContext );
            model.Material = material;
            model.IntersecitonProgPath = intersectPath;
            model.BoundingBoxProgPath = intersectPath;
            model.IntersecitonProgName = "mesh_intersect";
            model.BoundingBoxProgName = "mesh_bounds";
            model.GenerateNormals = true;
            model.ParseNormals = false;

            model.LoadContent();

            //build grid
            float maxDim = model.BBox.Extent( model.BBox.LongestAxis() );
            OptixV3 dims = model.BBox.Extent();
            OptixV3 voxelExtent = dims / maxDim * MaxVoxelSize;

            BuildGrid( model.BBox, -1, voxelExtent.LengthSqrd() );
            mRenderer.CreateVPLData( "../../Assets/Shaders/VPLMerge.hlsl", "CSMain", mVoxels, mFirstLeafIndex );

            SetCamera( model.BBox );

            /*-----------------------------------------------
             * Finally compile the optix context, and build the accel tree
             *-----------------------------------------------*/
            Console.WriteLine( "Compiling Optix scene." );
            mContext[ "top_object" ].Set( model.GeoGroup );

            mContext.Compile();
            mContext.BuildAccelTree();
        }

        public void Run()
        {
            MessagePump.Run( mForm, Render );

            foreach ( var item in SlimDX.ObjectTable.Objects )
                item.Dispose();
        }

        public void Render()
        {
            double start = OptixDotNet.Utils.Time.GetTimeInSecsDouble();

            mContext.Launch( 0, SqrtNumVPLs, SqrtNumVPLs );

            mRenderer.Render();

            double stop = OptixDotNet.Utils.Time.GetTimeInSecsDouble();
            mTimeDelta = (float)( stop - start );
        }

        private void SetCamera( BoundingBox box )
        {
            mCamera = new D3DCamera( 30.0f, (float)mForm.ClientSize.Width / (float)mForm.ClientSize.Height, 0.1f, 500.0f );
            mCamera.CenterOnBoundingBox( box, 0.7f );
            mCamera.RotationVel = 100.0f;
            mCamera.TranslationVel = 1000.0f;

            mRenderer.SetCamera( mCamera.CameraData );
        }

        private void BuildGrid( BoundingBox bounds, int index, float minExtent )
        {
            //if we're at the leaf node level, create a virtual point light for the cell
            if( bounds.Extent().LengthSqrd() <= minExtent )
            {
                mVoxels.Add( bounds );

                //this will be the first entry into the voxel list
                //we store the index of the cell so we can directly index the voxel array when we're searching the grid
                if ( mFirstLeafIndex == -1 )
                {
                    mFirstLeafIndex = index;
                }
                return;
            }

            OptixV3 dims = bounds.Extent() * .5f;

            //top left
            OptixV3 min = bounds.Min;
            BuildGrid( new BoundingBox( min, min + dims ), index * 8 + 1, minExtent );

            //top right
            min = bounds.Min + new OptixV3( dims.X, 0.0f, 0.0f );
            BuildGrid( new BoundingBox( min, min + dims ), index * 8 + 2, minExtent );

            //bottom right
            min = bounds.Min + new OptixV3( dims.X, 0.0f, dims.Z );
            BuildGrid( new BoundingBox( min, min + dims ), index * 8 + 3, minExtent );

            //bottom left
            min = bounds.Min + new OptixV3( 0.0f, 0.0f, dims.Z );
            BuildGrid( new BoundingBox( min, min + dims ), index * 8 + 4, minExtent );

            //top left
            min = bounds.Min + new OptixV3( 0.0f, dims.Y, 0.0f );
            BuildGrid( new BoundingBox( min, min + dims ), index * 8 + 5, minExtent );

            //top right
            min = bounds.Min + new OptixV3( dims.X, dims.Y, 0.0f );
            BuildGrid( new BoundingBox( min, min + dims ), index * 8 + 6, minExtent );

            //bottom right
            min = bounds.Min + new OptixV3( dims.X, dims.Y, dims.Z );
            BuildGrid( new BoundingBox( min, min + dims ), index * 8 + 7, minExtent );

            //bottom left
            min = bounds.Min + new OptixV3( 0.0f, dims.Y, dims.Z );
            BuildGrid( new BoundingBox( min, min + dims ), index * 8 + 8, minExtent );
        }

        private Vector3 OptixToSlimDX( OptixDotNet.Math.Vector3 vector )
        {
            return new Vector3( vector.X, vector.Y, vector.Z );
        }

        private void KeyboardInput( object sender, KeyboardInputEventArgs e )
        {
            float step = mTimeDelta * .01f;
            if ( e.State == KeyState.Pressed )
            {
                switch ( e.Key )
                {
                    case Keys.W:
                        mCamera.Dolly( step );
                        break;
                    case Keys.S:
                        mCamera.Dolly( -step );
                        break;
                    case Keys.A:
                        mCamera.Pan( -step, 0.0f );
                        break;
                    case Keys.D:
                        mCamera.Pan( step, 0.0f );
                        break;
                    default:
                        break;
                }

                mCamera.BuildView();
                mRenderer.SetCamera( mCamera.CameraData );
            }
            else
            {
                switch ( e.Key )
                {
                    case Keys.Escape:
                        mForm.Close();
                        Application.Exit();
                        break;
                    default:
                        break;
                }
            }
        }

        private void MouseInput( object sender, MouseInputEventArgs e )
        {
            float deltaX = e.X / (float)mForm.ClientSize.Width;
            float deltaY = e.Y / (float)mForm.ClientSize.Height;

            deltaX *= mTimeDelta;
            deltaY *= mTimeDelta;

            float deltaT = deltaX;
            if ( Math.Abs( deltaX ) < Math.Abs( deltaY ) )
                deltaT = deltaY;

            switch ( e.ButtonFlags )
            {
                case MouseButtonFlags.LeftDown:
                    mLeftMouseDown = true;
                    break;
                case MouseButtonFlags.LeftUp:
                    mLeftMouseDown = false;
                    break;
                case MouseButtonFlags.RightDown:
                    mRightMouseDown = true;
                    break;
                case MouseButtonFlags.RightUp:
                    mRightMouseDown = false;
                    break;
                case MouseButtonFlags.MiddleDown:
                    mMiddleMouseDown = true;
                    break;
                case MouseButtonFlags.MiddleUp:
                    mMiddleMouseDown = false;
                    break;
                default:
                    break;
            }

            if ( mLeftMouseDown )
            {
                mCamera.Rotate( -deltaX, -deltaY );
                mRenderer.SetCamera( mCamera.CameraData );
            }

            if ( mRightMouseDown )
            {
                mCamera.Dolly( -deltaT );
                mCamera.BuildView();
                mRenderer.SetCamera( mCamera.CameraData );
            }

            if ( mMiddleMouseDown )
            {
                mCamera.Pan( deltaX, -deltaY );
                mCamera.BuildView();
                mRenderer.SetCamera( mCamera.CameraData );
            }

            mPrevMouse = e;
        }
    }
}
